package com.dommettj.rawfix.session;

import com.dommettj.rawfix.message.MessageHandler;
import com.dommettj.rawfix.message.MessageType;
import com.dommettj.rawfix.message.MessageUtils;
import com.dommettj.rawfix.message.UnknownMessageTypeException;
import org.apache.log4j.Logger;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.nio.channels.spi.SelectorProvider;
import java.nio.charset.Charset;
import java.nio.charset.CharsetDecoder;
import java.util.*;
import java.util.concurrent.LinkedBlockingQueue;

public class NioSessionWorker implements Runnable {

    private static final Logger logger = Logger.getLogger(NioSessionWorker.class);
    private static final int READ_ONLY_SELECTION_KEY = SelectionKey.OP_CONNECT | SelectionKey.OP_READ;
    private static final int READ_WRITE_SELECTION_KEY = READ_ONLY_SELECTION_KEY | SelectionKey.OP_WRITE;
    private static Charset charset = Charset.forName("UTF-8");
    private static CharsetDecoder charsetDecoder = charset.newDecoder();
    private final ByteBuffer readBuffer = ByteBuffer.allocate(1024);
    private final ByteBuffer writeBuffer = ByteBuffer.allocate(1024);
    private final Queue<String> writeMessageQueue = new LinkedBlockingQueue<String>(16);
    private final Map<MessageType, Collection<MessageHandler>> messageHandlers;
    private final SocketChannel socketChannel;
    private final Selector selector;
    private volatile boolean shutdownRequested;

    public NioSessionWorker(final String hostName, final int port) throws IOException {

        messageHandlers = new HashMap<MessageType, Collection<MessageHandler>>();
        selector = SelectorProvider.provider().openSelector();
        socketChannel = SocketChannel.open();
        socketChannel.configureBlocking(false);
        socketChannel.connect(new InetSocketAddress(hostName, port));
        socketChannel.register(selector, READ_ONLY_SELECTION_KEY);
    }

    public void registerMessageHandler(final MessageType messageType, final MessageHandler messageHandler) {

        Collection<MessageHandler> handlers = messageHandlers.get(messageType);
        if(handlers == null) {
            handlers = new ArrayList<MessageHandler>();
            messageHandlers.put(messageType, handlers);
        }
        handlers.add(messageHandler);
    }

    public void sendMessage(final String message) throws ClosedChannelException {
        if(!writeMessageQueue.offer(message)) {
            logger.error("Failed to queue message: " + message);
        } else {
            socketChannel.register(selector, READ_WRITE_SELECTION_KEY);
            selector.wakeup();
        }
    }

    public void requestShutdown() {
        shutdownRequested = true;
    }

    @Override
    public void run() {

        shutdownRequested = false;

        while (!shutdownRequested) {

            try {
                selector.select();
            } catch (IOException e) {
                logger.error("Selector fail", e);
                break;
            }

            final Iterator<SelectionKey> selectionKeys = selector.selectedKeys().iterator();

            while (selectionKeys.hasNext()) {
                final SelectionKey key = selectionKeys.next();
                selectionKeys.remove();

                try {
                    processSelectionKey(key);
                } catch (IOException e) {
                    logger.error("Failed to process key: " + key, e);
                }
            }
        }
        try {
            selector.close();
        } catch (IOException e) {
            logger.error("Failed to close selector", e);
        }
    }

    private void processSelectionKey(final SelectionKey key) throws IOException {

        final SocketChannel channel = (SocketChannel) key.channel();

        if(key.isValid() && key.isConnectable()) {
            logger.debug("Connectable");
            if (!channel.finishConnect()) {
                logger.error("Failed to connect");
                key.cancel();
            }
        }

        if (key.isValid() && key.isReadable()) {
            logger.debug("Readable");
            read(channel);
        }

        if (key.isValid() && key.isWritable()) {

            logger.debug("Writable");
            if(writeMessageQueue.peek() != null) {
                writeMessage(channel, writeMessageQueue.poll());
            } else {
                socketChannel.register(selector, READ_ONLY_SELECTION_KEY);
            }
        }
    }

    private void read(final SocketChannel channel) {

        try {
            readBuffer.clear();
            int numBytesRead = channel.read(readBuffer);

            if (numBytesRead == -1) {
                channel.close();
            } else {
                readBuffer.flip();

                int oldPosition = readBuffer.position();
                final String content = charsetDecoder.decode(readBuffer).toString();

                // reset buffer's position to its original so it is not altered:
                readBuffer.position(oldPosition);

                logger.debug("Read " + numBytesRead + " bytes: " + content);
                handleIncomingContent(content);
            }
        } catch (IOException e) {
            logger.error("Failed to read", e);
        }
    }

    private void writeMessage(final SocketChannel channel, final String message) {

        try {
            writeBuffer.clear();
            final byte[] messageBytes = message.getBytes(charset);

            for(byte b : messageBytes) {
                writeBuffer.put(b);
            }

            writeBuffer.flip();
            int numBytesWritten = channel.write(writeBuffer);

            logger.debug("Wrote " + numBytesWritten + " bytes: " + message);

        } catch (IOException e) {
            logger.error("Failed to write", e);
        }
    }

    private void handleIncomingContent(final String content) {

        for(final String message : MessageUtils.getMessages(content)) {

            try {
                final MessageType messageType = MessageUtils.getMessageType(message);

                if(messageHandlers.get(messageType) != null) {
                    for(MessageHandler handler : messageHandlers.get(messageType)) {
                        handler.handle(message);
                    }
                } else {
                    logger.warn("No handlers registered for message type [" + messageType + "]");
                }
            } catch (UnknownMessageTypeException e) {
                logger.error("Cannot handle message [" + message + "]");
            }
        }
    }
}
